{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "n_var = 4\n",
    "embedding_dim = 24\n",
    "embedding_output = torch.randn((batch_size, n_var, embedding_dim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 24])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_attention_heads = 2\n",
    "attention_head_size = embedding_dim // 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_output.size()[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transpose_for_scores(x):\n",
    "    new_x_shape = x.size()[:-1] + (num_attention_heads, attention_head_size)\n",
    "    x = x.view(*new_x_shape)\n",
    "    return x.permute(0, 2, 1, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_x_shape = embedding_output.size()[:-1]+ (num_attention_heads, attention_head_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 2, 12])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_x_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_output = embedding_output.view(*new_x_shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 2, 12])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 2, 4, 12])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_output.permute(0, 2, 1,3).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = embedding_dim\n",
    "all_head_size = embedding_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = nn.Linear(hidden_size, all_head_size)\n",
    "key = nn.Linear(hidden_size, all_head_size)\n",
    "value = nn.Linear(hidden_size, all_head_size)\n",
    "hidden_states = embedding_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_output = torch.randn((batch_size, n_var, embedding_dim))\n",
    "hidden_states = embedding_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixed_query_layer = query(hidden_states)\n",
    "mixed_key_layer = key(hidden_states)\n",
    "mixed_value_layer = value(hidden_states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 24])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mixed_query_layer.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_layer = transpose_for_scores(mixed_query_layer)\n",
    "key_layer = transpose_for_scores(mixed_key_layer)\n",
    "value_layer = transpose_for_scores(mixed_value_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 2, 4, 12])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_layer.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 2, 12, 4])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_layer.transpose(-1, -2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 2, 4, 4])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attention_scores.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_scores = attention_scores / math.sqrt(attention_head_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_probs = nn.Softmax(dim=-1)(attention_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "context_layer = torch.matmul(attention_probs, value_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "context_layer = context_layer.permute(0, 2, 1, 3).contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 2, 12])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_layer.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_context_layer_shape = context_layer.size()[:-2] + (all_head_size,)\n",
    "context_layer = context_layer.view(*new_context_layer_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 24])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_layer.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, \"embedding_size\"):\n",
    "            raise ValueError(\n",
    "                \"The hidden size (%d) is not a multiple of the number of attention \"\n",
    "                \"heads (%d)\" % (config.hidden_size, config.num_attention_heads)\n",
    "            )\n",
    "        self.output_attentions = config.output_attentions\n",
    "\n",
    "        self.num_attention_heads = config.num_attention_heads\n",
    "        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)\n",
    "        self.all_head_size = self.num_attention_heads * self.attention_head_size\n",
    "\n",
    "        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.key = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.value = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "\n",
    "        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n",
    "\n",
    "    def transpose_for_scores(self, x):\n",
    "        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)\n",
    "        x = x.view(*new_x_shape)\n",
    "        return x.permute(0, 2, 1, 3)\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        mixed_query_layer = self.query(hidden_states)\n",
    "        mixed_key_layer = self.key(hidden_states)\n",
    "        mixed_value_layer = self.value(hidden_states)\n",
    "\n",
    "        query_layer = self.transpose_for_scores(mixed_query_layer)\n",
    "        key_layer = self.transpose_for_scores(mixed_key_layer)\n",
    "        value_layer = self.transpose_for_scores(mixed_value_layer)\n",
    "        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))\n",
    "        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n",
    "        attention_probs = nn.Softmax(dim=-1)(attention_scores)\n",
    "        attention_probs = self.dropout(attention_probs)\n",
    "        context_layer = torch.matmul(attention_probs, value_layer)\n",
    "\n",
    "        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n",
    "        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n",
    "        context_layer = context_layer.view(*new_context_layer_shape)\n",
    "\n",
    "        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerConfig:\n",
    "    def __init__(self, \n",
    "                output_attentions = True,\n",
    "                num_attention_heads = 2,\n",
    "                hidden_size = 24,\n",
    "                attention_probs_dropout_prob = 0.1,\n",
    "                hidden_dropout_prob = 0.1,\n",
    "                intermediate_size = 24):\n",
    "        self.output_attentions = output_attentions\n",
    "        self.num_attention_heads = num_attention_heads\n",
    "        self.hidden_size = hidden_size\n",
    "        self.attention_probs_dropout_prob = attention_probs_dropout_prob\n",
    "        self.hidden_dropout_prob = hidden_dropout_prob\n",
    "        self.intermediate_size = intermediate_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_config = TransformerConfig()\n",
    "self_attention = SelfAttention(transformer_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 24])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "self_attention(embedding_output)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerOutput(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
    "        self.LayerNorm = nn.LayerNorm(config.hidden_size)\n",
    "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
    "\n",
    "    def forward(self, hidden_states, input_tensor):\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = self.dropout(hidden_states)\n",
    "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
    "        return hidden_states\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerIntermediate(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)\n",
    "        self.intermediate_act_fn = F.relu\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = self.intermediate_act_fn(hidden_states)\n",
    "        return hidden_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerLayer(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.attention = SelfAttention(config)\n",
    "        self.output = TransformerOutput(config)\n",
    "        self.intermediate = TransformerIntermediate(config)\n",
    "    def forward(self, hidden):\n",
    "        self_attention_outputs = self.attention(hidden)\n",
    "        attention_output = self_attention_outputs[0]\n",
    "        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights\n",
    "\n",
    "        intermediate_output = self.intermediate(attention_output)\n",
    "        layer_output = self.output(intermediate_output, attention_output)\n",
    "        outputs = (layer_output,) + outputs\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_layer = TransformerLayer(transformer_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 4, 24])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformer_layer(embedding_output)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
